//
//  AgoraVideoFormatConverter.swift
//  AgoraPremium
//
//  Created by GongYuhua on 2017/11/29.
//  Copyright © 2017年 Agora. All rights reserved.
//

#if os(iOS)
    import UIKit
#else
    import Cocoa
#endif

struct AgoraVideoFormatConverter {
    static func convertPixelBufferFromNV12ToI420(_ buffer: CVPixelBuffer) -> CVPixelBuffer {
        CVPixelBufferLockBaseAddress(buffer, .readOnly)
        defer {
            CVPixelBufferUnlockBaseAddress(buffer, .readOnly)
        }
        
        let isPlanar = CVPixelBufferIsPlanar(buffer)
        let width = isPlanar ? CVPixelBufferGetWidthOfPlane(buffer, 0) : CVPixelBufferGetWidth(buffer)
        let height = isPlanar ? CVPixelBufferGetHeightOfPlane(buffer, 0) : CVPixelBufferGetHeight(buffer)
        
        let stride0 = CVPixelBufferGetBytesPerRowOfPlane(buffer, 0)
        let stride1 = CVPixelBufferGetBytesPerRowOfPlane(buffer, 1)
        
        let src0 = CVPixelBufferGetBaseAddressOfPlane(buffer, 0)!.assumingMemoryBound(to: uint8.self)
        let src1 = CVPixelBufferGetBaseAddressOfPlane(buffer, 1)!.assumingMemoryBound(to: uint8.self)
        
        let attributes:[NSString: NSDictionary] = [kCVPixelBufferIOSurfacePropertiesKey: [:]]
        var pixelBuffer: CVPixelBuffer?
        CVPixelBufferCreate(kCFAllocatorDefault, width, height, kCVPixelFormatType_420YpCbCr8PlanarFullRange, attributes as CFDictionary, &pixelBuffer)
        
        let newBuffer = pixelBuffer!
        CVPixelBufferLockBaseAddress(newBuffer, .readOnly)
        defer {
            CVPixelBufferUnlockBaseAddress(newBuffer, .readOnly)
        }
        
        let dstY = CVPixelBufferGetBaseAddressOfPlane(newBuffer, 0)!.assumingMemoryBound(to: uint8.self)
        let strideY = CVPixelBufferGetBytesPerRowOfPlane(newBuffer, 0)
        let dstU = CVPixelBufferGetBaseAddressOfPlane(newBuffer, 1)!.assumingMemoryBound(to: uint8.self)
        let strideU = CVPixelBufferGetBytesPerRowOfPlane(newBuffer, 1)
        let dstV = CVPixelBufferGetBaseAddressOfPlane(newBuffer, 2)!.assumingMemoryBound(to: uint8.self)
        let strideV = CVPixelBufferGetBytesPerRowOfPlane(newBuffer, 2)
        NV12ToI420(src0, Int32(stride0),
                   src1, Int32(stride1),
                   dstY, Int32(strideY),
                   dstU, Int32(strideU),
                   dstV, Int32(strideV),
                   Int32(width), Int32(height))
        
        return newBuffer
    }
    
    
    
    static func convertPixelBufferFromI420(_ buffer: NSData,imageSize width:Int,height:Int ) -> CVPixelBuffer {

          let attributes:[NSString: NSDictionary] = [kCVPixelBufferIOSurfacePropertiesKey: [:]]
          var pixelBuffer: CVPixelBuffer?
          CVPixelBufferCreate(kCFAllocatorDefault, width, height, kCVPixelFormatType_420YpCbCr8PlanarFullRange, attributes as CFDictionary, &pixelBuffer)

          let newBuffer = pixelBuffer!
          CVPixelBufferLockBaseAddress(newBuffer, .readOnly)
          defer {
              CVPixelBufferUnlockBaseAddress(newBuffer, .readOnly)
          }
          let srcY = buffer.bytes;
          let srcU = buffer.bytes+width*height;
          let srcV = buffer.bytes+width*height*5/4;
          let srcStrideY = width;
          let srcStrideU = width/2;
          let srcStrideV = width/2;

          let dstY = CVPixelBufferGetBaseAddressOfPlane(newBuffer, 0)!.assumingMemoryBound(to: uint8.self)
          let strideY = CVPixelBufferGetBytesPerRowOfPlane(newBuffer, 0)
          let dstU = CVPixelBufferGetBaseAddressOfPlane(newBuffer, 1)!.assumingMemoryBound(to: uint8.self)
          let strideU = CVPixelBufferGetBytesPerRowOfPlane(newBuffer, 1)
          let dstV = CVPixelBufferGetBaseAddressOfPlane(newBuffer, 2)!.assumingMemoryBound(to: uint8.self)
          let strideV = CVPixelBufferGetBytesPerRowOfPlane(newBuffer, 2)
         
          for i in 0 ..< height{
             memcpy(dstY+strideY*i, srcY+srcStrideY*i, srcStrideY);
          }
          for i in 0 ..< height/2{
             memcpy(dstU+strideU*i, srcU+srcStrideU*i, srcStrideU);
          }
          for i in 0 ..< height/2{
             memcpy(dstV+strideV*i, srcV+srcStrideV*i, srcStrideV);
          }
                 
         
          return pixelBuffer!
      }
    
    static func convertPixelBuffer(_ buffer: CVPixelBuffer, toRawDataOf format: AgoraVideoPixelFormat) -> (pointer: UnsafeMutableRawPointer, size: CGSize) {
        CVPixelBufferLockBaseAddress(buffer, .readOnly)
        defer {
            CVPixelBufferUnlockBaseAddress(buffer, .readOnly)
        }
        
        let isPlanar = CVPixelBufferIsPlanar(buffer)
        let width = isPlanar ? CVPixelBufferGetWidthOfPlane(buffer, 0) : CVPixelBufferGetWidth(buffer)
        let height = isPlanar ? CVPixelBufferGetHeightOfPlane(buffer, 0) : CVPixelBufferGetHeight(buffer)
        let size = CGSize(width: width, height: height)
        let area = width * height
        
        let result: (UnsafeMutableRawPointer, CGSize)
        switch format {
        case .BGRA:
            // Source will be BGRA
            let ybuffer = CVPixelBufferGetBaseAddressOfPlane(buffer, 0)!
            result = (ybuffer, size)
            
        case .NV12:
            // Source will be NV12
            let src0 = CVPixelBufferGetBaseAddressOfPlane(buffer, 0)!.assumingMemoryBound(to: uint8.self)
            let src1 = CVPixelBufferGetBaseAddressOfPlane(buffer, 1)!.assumingMemoryBound(to: uint8.self)
            
            let dst = UnsafeMutablePointer<uint8>.allocate(capacity: area * 3 / 2)
            let pointer = UnsafeMutableRawPointer(dst)
            memcpy(pointer, src0, area)
            memcpy(pointer + area, src1, area / 2)
            result = (UnsafeMutableRawPointer(dst), size)
            
        case .I420:
            // Source will be NV12
            let src0 = CVPixelBufferGetBaseAddressOfPlane(buffer, 0)!.assumingMemoryBound(to: uint8.self)
            let src1 = CVPixelBufferGetBaseAddressOfPlane(buffer, 1)!.assumingMemoryBound(to: uint8.self)
            let stride0 = CVPixelBufferGetBytesPerRowOfPlane(buffer, 0)
            let stride1 = CVPixelBufferGetBytesPerRowOfPlane(buffer, 1)
            
            let dst = UnsafeMutablePointer<uint8>.allocate(capacity: area * 3 / 2)
            NV12ToI420(src0, Int32(stride0), src1, Int32(stride1),
                       dst, Int32(width),
                       dst + area, Int32(width / 2),
                       dst + area * 5 / 4, Int32(width / 2),
                       Int32(width), Int32(height))
            
            result = (UnsafeMutableRawPointer(dst), size)
        }
        
        return result
    }
}
